import xarray as xr
import numpy as np
import os
from datetime import datetime, timedelta
import re

def parse_filename_time(filename):
    """
    从文件名中解析起报时间和预报时效
    
    文件名格式: plev_linjin_2025062200_000.nc
    - 2025062200: 起报时间 (YYYYMMDDHH)
    - 000: 预报时效 (HHH小时)
    
    Returns:
    --------
    init_time_str : str
        起报时间字符串，格式: "2025-06-22 00:00:00"
    forecast_hour : int
        预报时效（小时）
    valid_time_str : str
        有效时间字符串，格式: "2025-06-22 00:00:00"
    """
    
    basename = os.path.basename(filename)
    
    # 使用正则表达式提取时间信息
    # 匹配格式: plev_xxx_YYYYMMDDHH_HHH.nc
    pattern = r'plev_\w+_(\d{10})_(\d{3})\.nc'
    match = re.search(pattern, basename)
    
    if not match:
        raise ValueError(f"无法从文件名解析时间信息: {basename}")
    
    init_time_str = match.group(1)  # 2025062200
    lead_time_str = match.group(2)  # 000
    
    # 解析起报时间
    year = int(init_time_str[:4])
    month = int(init_time_str[4:6])
    day = int(init_time_str[6:8])
    hour = int(init_time_str[8:10])
    
    init_datetime = datetime(year, month, day, hour)
    forecast_hour = int(lead_time_str)
    
    # 计算有效时间
    valid_datetime = init_datetime + timedelta(hours=forecast_hour)
    
    # 格式化时间字符串
    init_time_str = init_datetime.strftime("%Y-%m-%d %H:%M:%S")
    valid_time_str = valid_datetime.strftime("%Y-%m-%d %H:%M:%S")
    
    print(f"从文件名解析的时间信息:")
    print(f"  起报时间: {init_time_str}")
    print(f"  预报时效: {forecast_hour} 小时")
    print(f"  有效时间: {valid_time_str}")
    
    return init_time_str, forecast_hour, valid_time_str

def convert_to_met_cf_compliant(input_file, output_file):
    """
    将NetCDF文件转换为MET要求的CF-compliant格式
    从文件名解析时间信息
    
    Parameters:
    -----------
    input_file : str
        输入NetCDF文件路径
    output_file : str  
        输出CF-compliant NetCDF文件路径
    """
    
    print(f"正在转换文件: {input_file}")
    
    # 从文件名解析时间信息
    try:
        init_time_str, forecast_hour, valid_time_str = parse_filename_time(input_file)
    except ValueError as e:
        print(f"警告: {e}")
        # 如果解析失败，使用默认值
        init_time_str = "2025-06-22 00:00:00"
        forecast_hour = 0
        valid_time_str = "2025-06-22 00:00:00"
    
    # 读取原始文件
    ds = xr.open_dataset(input_file)
    
    # ========== 1. 添加必需的全局属性 ==========
    ds.attrs['Conventions'] = 'CF-1.6'  # 必需：CF约定声明
    
    # 保留有用的描述信息，清理WRF特有的属性
    ds.attrs['title'] = 'CESM model output on pressure levels'
    ds.attrs['institution'] = 'Your Institution'
    ds.attrs['source'] = 'CESM/WRF model'
    if 'description' in ds.attrs:
        ds.attrs['description'] = ds.attrs['description']
    
    # 移除WRF特有的全局属性
    wrf_attrs = ['MAP_PROJ', 'MAP_PROJ_CHAR', 'TRUELAT1', 'TRUELAT2', 
                 'STAND_LON', 'MOAD_CEN_LAT', 'CEN_LAT', 'POLE_LAT', 
                 'POLE_LON', 'DX', 'DY']
    for attr in wrf_attrs:
        if attr in ds.attrs:
            del ds.attrs[attr]
    
    # ========== 2. 修改坐标变量属性 ==========
    
    # 时间坐标 - 根据文件名设置正确的时间信息
    time_units = f"hours since {init_time_str}"
    
    # 设置time变量的值为预报时效
    ds['time'] = xr.DataArray([forecast_hour], dims=['time'], coords={'time': [forecast_hour]})
    
    ds['time'].attrs.update({
        'standard_name': 'time',           # CF标准名称
        'units': time_units,               # 从文件名解析的起报时间
        'axis': 'T',                       # 时间轴标识
        'long_name': 'time'
    })
    
    # 压强坐标
    ds['plev'].attrs.update({
        'standard_name': 'air_pressure',   # CF标准名称
        'axis': 'Z',                       # 垂直轴标识
        'long_name': 'pressure level'
    })
    # 保持原有的units='hPa'和positive='down'
    
    # 纬度坐标 - 添加必需的units
    ds['lat'].attrs.update({
        'standard_name': 'latitude',       # CF标准名称
        'units': 'degrees_north',          # 必需的单位
        'axis': 'Y',                       # 纬度轴标识
        'long_name': 'latitude'
    })
    # 移除_FillValue（坐标变量不应该有）
    if '_FillValue' in ds['lat'].attrs:
        del ds['lat'].attrs['_FillValue']
    
    # 经度坐标 - 添加必需的units  
    ds['lon'].attrs.update({
        'standard_name': 'longitude',      # CF标准名称
        'units': 'degrees_east',           # 必需的单位
        'axis': 'X',                       # 经度轴标识
        'long_name': 'longitude'
    })
    # 移除_FillValue（坐标变量不应该有）
    if '_FillValue' in ds['lon'].attrs:
        del ds['lon'].attrs['_FillValue']
    
    # ========== 3. 修改数据变量属性 ==========
    
    # 温度 - 转换单位从摄氏度到开尔文
    if ds['Temperature'].attrs.get('units') == 'degC':
        print("转换温度单位：摄氏度 -> 开尔文")
        ds['Temperature'] = ds['Temperature'] + 273.15
    
    ds['Temperature'].attrs.update({
        'standard_name': 'air_temperature',
        'units': 'K',
        'long_name': 'Air temperature',
        'coordinates': 'time plev lat lon'  # 指定坐标变量
    })
    
    # 东风分量
    ds['UWind'].attrs.update({
        'standard_name': 'eastward_wind',
        'units': 'm s-1',
        'long_name': 'Eastward wind component',
        'coordinates': 'time plev lat lon'
    })
    
    # 北风分量
    ds['VWind'].attrs.update({
        'standard_name': 'northward_wind',
        'units': 'm s-1', 
        'long_name': 'Northward wind component',
        'coordinates': 'time plev lat lon'
    })
    
    # 比湿
    ds['SpecificHumidity'].attrs.update({
        'standard_name': 'specific_humidity',
        'units': 'kg kg-1',
        'long_name': 'Specific humidity',
        'coordinates': 'time plev lat lon'
    })
    
    # 位势高度
    ds['GeopotentialHeight'].attrs.update({
        'standard_name': 'geopotential_height',
        'units': 'm',
        'long_name': 'Geopotential height',
        'coordinates': 'time plev lat lon'
    })
    
    # 密度
    ds['Density'].attrs.update({
        'standard_name': 'air_density',
        'units': 'kg m-3',
        'long_name': 'Air density',
        'coordinates': 'time plev lat lon'
    })
    
    # ========== 4. 添加预报参考时间 ==========
    # 设置forecast_reference_time为起报时间（时效0）
    ds['forecast_reference_time'] = xr.DataArray(
        data=0.0,  # 起报时间对应时效0
        attrs={
            'standard_name': 'forecast_reference_time',
            'units': time_units,  # 使用相同的时间基准
            'long_name': 'forecast reference time'
        }
    )
    
    # ========== 5. 数据类型优化 ==========
    # 设置编码以优化文件大小和兼容性
    encoding = {}
    
    # 坐标变量保持双精度
    encoding['time'] = {'dtype': 'float64'}
    encoding['plev'] = {'dtype': 'float64'}
    encoding['lat'] = {'dtype': 'float64'}
    encoding['lon'] = {'dtype': 'float64'}
    encoding['forecast_reference_time'] = {'dtype': 'float64'}
    
    # 数据变量使用单精度节省空间
    for var in ['Temperature', 'SpecificHumidity', 'UWind', 'VWind', 
                'GeopotentialHeight', 'Density']:
        if var in ds.data_vars:
            encoding[var] = {
                'dtype': 'float32',
                'zlib': True,           # 启用压缩
                'complevel': 4,         # 压缩级别
                'shuffle': True         # 启用shuffle filter
            }
    
    # ========== 6. 保存文件 ==========
    print(f"正在保存CF-compliant文件: {output_file}")
    ds.to_netcdf(output_file, encoding=encoding)
    ds.close()
    
    print("转换完成！主要修改:")
    print("✓ 添加了 Conventions='CF-1.6' 全局属性")
    print("✓ 从文件名解析并设置了正确的时间信息")
    print("✓ 为所有坐标变量添加了 standard_name 和 units 属性")
    print("✓ 为所有数据变量添加了 standard_name 和 coordinates 属性") 
    print("✓ 温度单位从摄氏度转换为开尔文")
    print("✓ 添加了 forecast_reference_time 变量")
    print("✓ 优化了数据类型和压缩设置")
    
    return output_file

def batch_convert_with_filename_parsing(input_files, output_dir=None):
    """
    批量转换多个NetCDF文件，自动从文件名解析时间
    
    Parameters:
    -----------
    input_files : list
        输入文件路径列表
    output_dir : str, optional
        输出目录，如果不指定则在原文件目录生成
    """
    
    for input_file in input_files:
        if output_dir:
            filename = os.path.basename(input_file)
            name, ext = os.path.splitext(filename)
            output_file = os.path.join(output_dir, f"{name}_cf{ext}")
        else:
            name, ext = os.path.splitext(input_file)
            output_file = f"{name}_cf{ext}"
        
        try:
            convert_to_met_cf_compliant(input_file, output_file)
            print(f"✓ 成功转换: {input_file} -> {output_file}")
        except Exception as e:
            print(f"✗ 转换失败 {input_file}: {e}")
            import traceback
            traceback.print_exc()

def verify_cf_compliance(file_path):
    """
    验证转换后的文件是否符合MET CF-compliant要求
    """
    
    print(f"\n验证文件: {file_path}")
    print("="*50)
    
    ds = xr.open_dataset(file_path)
    
    # 检查必需的全局属性
    print("1. 检查全局属性:")
    if 'Conventions' in ds.attrs:
        print(f"   ✓ Conventions: {ds.attrs['Conventions']}")
    else:
        print("   ✗ 缺少 Conventions 属性")
    
    # 检查坐标变量
    print("\n2. 检查坐标变量:")
    coord_vars = ['time', 'plev', 'lat', 'lon']
    
    for coord in coord_vars:
        if coord in ds.coords:
            attrs = ds[coord].attrs
            values = ds[coord].values
            print(f"   {coord}:")
            print(f"      standard_name: {attrs.get('standard_name', '缺失')}")
            print(f"      units: {attrs.get('units', '缺失')}")
            print(f"      axis: {attrs.get('axis', '缺失')}")
            if coord == 'time':
                print(f"      值: {values}")
        else:
            print(f"   ✗ 缺少坐标变量: {coord}")
    
    # 检查预报参考时间
    print("\n3. 检查预报参考时间:")
    if 'forecast_reference_time' in ds.variables:
        frt = ds['forecast_reference_time']
        print(f"   ✓ forecast_reference_time: {frt.values}")
        print(f"      units: {frt.attrs.get('units', '缺失')}")
        print(f"      standard_name: {frt.attrs.get('standard_name', '缺失')}")
    else:
        print("   - 未包含 forecast_reference_time")
    
    # 检查数据变量
    print("\n4. 检查数据变量:")
    data_vars = ['Temperature', 'UWind', 'VWind', 'SpecificHumidity']
    
    for var in data_vars:
        if var in ds.data_vars:
            attrs = ds[var].attrs
            print(f"   {var}:")
            print(f"      standard_name: {attrs.get('standard_name', '缺失')}")
            print(f"      units: {attrs.get('units', '缺失')}")
            print(f"      coordinates: {attrs.get('coordinates', '缺失')}")
        else:
            print(f"   ✗ 缺少数据变量: {var}")
    
    ds.close()
    print("\n验证完成!")

# 使用示例
if __name__ == "__main__":
    # 单文件转换
    input_file = "plev_linjin_2025062200_000.nc"
    output_file = "plev_linjin_2025062200_000_cf.nc"
    try:
        # 执行转换
        convert_to_met_cf_compliant(input_file, output_file)
        # 验证结果
        verify_cf_compliance(output_file)
        print(f"\n转换成功！文件已保存为: {output_file}")
        print("可以在MET中使用以下配置:")
        print("""
fcst = {
    file_type = NETCDF_NCCF;
    field = [
      { name = "Temperature"; level = [ "850" ]; },
      { name = "UWind"; level = [ "850" ]; },
      { name = "VWind"; level = [ "850" ]; }
    ];
}
        """)
        
    except Exception as e:
        print(f"转换过程中出现错误: {e}")
        import traceback
        traceback.print_exc()
    
    # 批量转换示例
    # input_files = [
    #     "plev_linjin_2025062200_000.nc",
    #     "plev_linjin_2025062206_000.nc",
    #     "plev_linjin_2025062212_000.nc"
    # ]
    # batch_convert_with_filename_parsing(input_files)